from hashlib import md5
from PIL import Image
import numpy as np
import cv2
import pickle
import os
import io
import sys
from tqdm import tqdm

# sys.path.append("/home/swls/work_dir/git/python_script")

# from tool import darknet_tool
# from tool import via_tool # export PYTHONPATH=$PYTHONPATH:`pwd`

from tool import filesystem

def main(img_dir):
    # if os.path.isdir(img_dir):
    #     imgs = [img_dir + os.sep + f for f in os.listdir(img_dir)]
    # else:
    #     imgs = [img_dir]
    imgs = filesystem.get_all_filepath(img_dir, ["jpg"])

    # md5_dict = dict()
    remove_count = 0
    for idx, p in tqdm(enumerate(imgs), total=len(imgs)):
        # label = os.path.basename(p).split("_")[0]

        ## 1
        file_size = os.path.getsize(p)
        # print(file_size)
        if file_size < 50000:
            os.remove(p)
            
        ## 2
        # try:
        #     cv_image = cv2.imread(p)
        #     if cv_image is None :
        #         continue
        #     crop_image = cv_image[0:900, 0:1250]
        #     m = np.mean(crop_image)
        #     # print(m)

        #     if m < 5:
        #         os.remove(p)
        #         remove_count += 1
        # except:
        #     pass


        ## 3
        # pil_image = Image.open(p)
        # crop_image = pil_image.crop((0, 457, 923, 949))
        # # crop_image.show()
        # bytes_data = io.BytesIO()
        # crop_image.save(bytes_data, format="png")

        # hash_algorithm = md5()
        # hash_algorithm.update(bytes_data.getvalue())
        # md5_str = hash_algorithm.hexdigest()
        # print(idx, "   " ,md5_str)

            # md5_dict[md5_str] = label
    # print(len(md5_dict.keys()))
    # with open("data.pkl", "wb") as wf:
    #     pickle.dump(md5_dict, wf)

if __name__ == "__main__":
    img_dir = sys.argv[1]
    main(img_dir)